#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import torch
import torchaudio
from flask import Flask, request, send_file
from io import BytesIO
from cosyvoice.cli.cosyvoice import CosyVoice2
from cosyvoice.utils.file_utils import load_wav
import os
import re
import logging
import numpy as np

# 配置参数
MODEL_DIR = 'pretrained_models/CosyVoice2-0.5B'
REFERENCE_AUDIO_PATH = './asset/ccc.wav'
REFERENCE_TEXT = "9月4日，中共中央总书记"
HOST = '0.0.0.0'
PORT = 5001
SAMPLE_RATE = 24000

# 初始化日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger('TTS-Service')

# 全局模型实例
cosyvoice = None
reference_audio = None


def init_tts_system():
    """初始化文本转语音系统"""
    global cosyvoice, reference_audio

    if cosyvoice is not None and reference_audio is not None:
        return True

    logger.info("正在加载文本转语音模型...")

    try:
        # 加载CosyVoice2模型
        cosyvoice = CosyVoice2(
            MODEL_DIR,
            load_jit=False,
            load_trt=False,
            load_vllm=False,
            fp16=True  # 使用半精度加速推理
        )

        # 加载参考音频
        reference_audio = load_wav(REFERENCE_AUDIO_PATH, 16000)
        logger.info(f"模型加载完成，参考音频: {REFERENCE_AUDIO_PATH}")
        return True
    except Exception as e:
        logger.error(f"模型加载失败: {str(e)}")
        return False


def clean_text(text):
    """清理文本，确保只包含模型支持的字符"""
    # 移除非中文字符、标点等
    return re.sub(r'[^\u4e00-\u9fa5，。！？；：、\s\w]', '', text)


def text_to_speech(text):
    """核心文本转语音函数"""
    # 清理文本
    cleaned_text = clean_text(text)

    # 生成语音
    results = list(cosyvoice.inference_zero_shot(
        cleaned_text,
        REFERENCE_TEXT,
        reference_audio,
        stream=False,
        text_frontend=False
    ))

    if not results or 'tts_speech' not in results[0]:
        raise ValueError("语音生成失败")

    # 获取语音数据
    speech_data = results[0]['tts_speech']

    # 修正张量维度
    if speech_data.dim() == 1:
        speech_data = speech_data.unsqueeze(0)
    elif speech_data.dim() == 3:
        speech_data = speech_data.squeeze(0)

    # 确保是2D张量
    if speech_data.dim() != 2:
        raise ValueError(f"无效的音频张量维度: {speech_data.dim()}")

    return speech_data


# 初始化Flask应用
app = Flask(__name__)


@app.route('/synthesize', methods=['POST'])
def synthesize():
    """语音合成接口"""
    try:
        # 获取文本
        data = request.json
        if not data or 'text' not in data:
            return {"error": "缺少文本参数"}, 400

        text = data['text']
        if not text.strip():
            return {"error": "文本内容为空"}, 400

        logger.info(f"合成请求: {text[:50]}{'...' if len(text) > 50 else ''}")

        # 生成语音
        speech_data = text_to_speech(text)

        # 创建内存中的WAV文件
        buffer = BytesIO()
        torchaudio.save(buffer, speech_data, SAMPLE_RATE, format='wav')
        buffer.seek(0)

        logger.info(f"合成成功, 时长: {speech_data.size(1) / SAMPLE_RATE:.2f}秒")
        return send_file(
            buffer,
            mimetype='audio/wav',
            as_attachment=False
        )
    except Exception as e:
        logger.error(f"合成失败: {str(e)}")
        return {"error": str(e)}, 500


@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return {"status": "ok", "model_loaded": cosyvoice is not None}, 200


if __name__ == '__main__':
    # 初始化模型
    if not init_tts_system():
        logger.error("服务启动失败: 模型初始化失败")
        exit(1)

    # 预热模型
    try:
        logger.info("预热模型...")
        warmup_text = "预热测试"
        speech_data = text_to_speech(warmup_text)
        logger.info(f"预热成功, 生成音频时长: {speech_data.size(1) / SAMPLE_RATE:.2f}秒")
    except Exception as e:
        logger.error(f"预热失败: {str(e)}")

    # 启动服务
    logger.info(f"服务启动: http://{HOST}:{PORT}")
    app.run(host=HOST, port=PORT, threaded=True)